"""项目主入口"""
import argparse
import logging
from pathlib import Path
from dotenv import load_dotenv
from config.app_config import AppConfig
from src.document_processing.document_loader import DocumentLoaderFactory
from src.document_processing.text_splitter import TextSplitterFactory
from src.embedding.huggingface_embedding import HuggingFaceEmbedding
from src.vector_db.chroma_db import ChromaDB
from src.llm.llm_factory import LLMFactory
from src.qa_chain.retrieval_qa import RetrievalQAChain
from src.utils.logger import setup_logger

load_dotenv()

def setup_directories(config: AppConfig):
    directories = [
        config.knowledge_base_path,
        config.vector_db_path,
        config.models_path,
        config.logs_path,
        f"{config.knowledge_base_path}/pdfs",
        f"{config.models_path}/embeddings",
        f"{config.models_path}/llms"
    ]
    for dir_path in directories:
        Path(dir_path).mkdir(parents=True, exist_ok=True)

def initialize_components(config: AppConfig):
    # 初始化嵌入模型
    embedding = HuggingFaceEmbedding(
        model_name=config.embedding_config.model_name,
        model_path=config.embedding_config.model_path,
        device=config.embedding_config.device
    )

    # 初始化向量数据库
    vector_db = ChromaDB.load(
        embedding=embedding,
        db_path=config.vector_db_path,
        collection_name=config.collection_name
    )

    # 初始化LLM
    llm = LLMFactory.create_llm(config.llm_config)

    # 初始化问答链
    qa_chain = RetrievalQAChain(vector_db=vector_db, llm=llm)

    return {"vector_db": vector_db, "qa_chain": qa_chain}

def load_and_process_documents(config: AppConfig, vector_db):
    knowledge_base = Path(config.knowledge_base_path)
    document_paths = []
    for ext in config.supported_document_extensions:
        document_paths.extend(knowledge_base.rglob(f"*{ext}"))

    if not document_paths:
        logging.warning("未找到任何文档，请添加文件到knowledge_base目录")
        return

    total_chunks = 0
    for doc_path in document_paths:
        try:
            loader = DocumentLoaderFactory.get_loader(str(doc_path))
            documents = loader.load()

            doc_type = "markdown" if doc_path.suffix.lower() == ".md" else "default"
            splitter = TextSplitterFactory.get_splitter(
                doc_type=doc_type,
                chunk_size=config.chunk_size,
                chunk_overlap=config.chunk_overlap
            )
            chunks = splitter.split_documents(documents)

            if chunks:
                vector_db.add_documents(chunks)
                total_chunks += len(chunks)
                logging.info(f"处理完成：{doc_path}（{len(chunks)}块）")
        except Exception as e:
            logging.error(f"处理{doc_path}失败：{str(e)}")

    vector_db.persist()
    logging.info(f"共处理{total_chunks}个文本块")

def interactive_qa(qa_chain):
    print("===== zbs_agent 问答系统 =====")
    print("输入问题查询，'exit'退出，'clear'清屏")

    while True:
        question = input("\n请输入问题：").strip()
        if question.lower() == "exit":
            print("再见！")
            break
        if question.lower() == "clear":
            print("\033c", end="")
            continue
        if not question:
            continue

        try:
            result = qa_chain.run(question=question, k=5, return_sources=True)
            print("\n===== 回答 =====")
            print(result["answer"])

            if result.get("sources"):
                print("\n===== 参考来源 =====")
                for i, src in enumerate(result["sources"], 1):
                    print(f"{i}. {src['source']}（相关性：{src['score']:.4f}）")
        except Exception as e:
            print(f"处理失败：{str(e)}")

def main():
    parser = argparse.ArgumentParser(description="zbs_agent 知识库问答系统")
    parser.add_argument("--reload", action="store_true", help="重新加载文档")
    parser.add_argument("--debug", action="store_true", help="开启调试模式")
    args = parser.parse_args()

    config = AppConfig()
    setup_logger(config.logs_path, args.debug)
    setup_directories(config)

    components = initialize_components(config)
    vector_db = components["vector_db"]
    qa_chain = components["qa_chain"]

    if args.reload:
        load_and_process_documents(config, vector_db)

    interactive_qa(qa_chain)
    vector_db.close()

if __name__ == "__main__":
    main()